"""
检索策略模块

定义不同的检索策略和配置选项。
"""

from typing import Dict, List, Optional, Any
from enum import Enum
from dataclasses import dataclass


class RetrievalMode(Enum):
    """检索模式枚举"""
    STANDARD = "standard"      # 标准检索
    SIMILARITY = "similarity"  # 相似度检索
    HYBRID = "hybrid"         # 混合检索
    SEMANTIC = "semantic"     # 语义检索
    KEYWORD = "keyword"       # 关键词检索


class ResponseMode(Enum):
    """响应模式枚举"""
    COMPACT = "compact"           # 紧凑模式
    REFINE = "refine"            # 精炼模式
    TREE_SUMMARIZE = "tree_summarize"  # 树状摘要模式
    SIMPLE = "simple"             # 简单模式
    NO_TEXT = "no_text"           # 无文本模式
    GENERATE = "generate"         # 生成模式


@dataclass
class RetrievalConfig:
    """检索配置"""
    # 基本配置
    similarity_top_k: int = 5
    similarity_cutoff: Optional[float] = None
    response_mode: ResponseMode = ResponseMode.COMPACT

    # 高级配置
    enable_node_keywords: bool = True
    enable_sparse_embedding: bool = False
    vector_store_query_mode: str = "default"  # "default", "mmr", "simple"

    # 查询转换
    query_transform_mode: str = "default"  # "default", "hyde", "embedding"

    # 后处理
    node_postprocessors: List[str] = None  # ["SimilarityPostprocessor", "KeywordPostprocessor"]

    # 其他参数
    verbose: bool = False

    def __post_init__(self):
        """初始化后处理"""
        if self.node_postprocessors is None:
            self.node_postprocessors = []

    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        return {
            "similarity_top_k": self.similarity_top_k,
            "similarity_cutoff": self.similarity_cutoff,
            "response_mode": self.response_mode.value,
            "enable_node_keywords": self.enable_node_keywords,
            "enable_sparse_embedding": self.enable_sparse_embedding,
            "vector_store_query_mode": self.vector_store_query_mode,
            "query_transform_mode": self.query_transform_mode,
            "node_postprocessors": self.node_postprocessors,
            "verbose": self.verbose
        }

    @classmethod
    def from_dict(cls, config_dict: Dict[str, Any]) -> 'RetrievalConfig':
        """从字典创建配置"""
        # 处理枚举类型
        if 'response_mode' in config_dict:
            if isinstance(config_dict['response_mode'], str):
                config_dict['response_mode'] = ResponseMode(config_dict['response_mode'])

        return cls(**config_dict)


class RetrievalStrategy:
    """检索策略基类"""

    def __init__(self, config: RetrievalConfig):
        """
        初始化检索策略

        Args:
            config: 检索配置
        """
        self.config = config

    def get_retriever_kwargs(self) -> Dict[str, Any]:
        """
        获取检索器参数

        Returns:
            Dict[str, Any]: 检索器参数字典
        """
        return {
            "similarity_top_k": self.config.similarity_top_k,
            "vector_store_query_mode": self.config.vector_store_query_mode
        }

    def get_query_engine_kwargs(self) -> Dict[str, Any]:
        """
        获取查询引擎参数

        Returns:
            Dict[str, Any]: 查询引擎参数字典
        """
        return {
            "response_mode": self.config.response_mode.value,
            "verbose": self.config.verbose,
            "node_postprocessors": self.config.node_postprocessors,
            "similarity_top_k": self.config.similarity_top_k
        }

    def validate_query(self, query: str) -> bool:
        """
        验证查询是否适合此策略

        Args:
            query: 查询文本

        Returns:
            bool: 是否适合
        """
        return len(query.strip()) > 0


class StandardRetrievalStrategy(RetrievalStrategy):
    """标准检索策略"""

    def __init__(self, config: Optional[RetrievalConfig] = None):
        super().__init__(config or RetrievalConfig())


class SimilarityRetrievalStrategy(RetrievalStrategy):
    """相似度检索策略"""

    def __init__(self, similarity_cutoff: float = 0.7, top_k: int = 10):
        config = RetrievalConfig(
            similarity_cutoff=similarity_cutoff,
            similarity_top_k=top_k,
            vector_store_query_mode="default"
        )
        super().__init__(config)


class HybridRetrievalStrategy(RetrievalStrategy):
    """混合检索策略（向量+关键词）"""

    def __init__(self, top_k: int = 5, sparse_weight: float = 0.5):
        config = RetrievalConfig(
            similarity_top_k=top_k,
            enable_sparse_embedding=True,
            vector_store_query_mode="hybrid",
            node_postprocessors=["SimilarityPostprocessor"]
        )
        super().__init__(config)
        self.sparse_weight = sparse_weight

    def get_retriever_kwargs(self) -> Dict[str, Any]:
        kwargs = super().get_retriever_kwargs()
        kwargs.update({
            "sparse_top_k": self.config.similarity_top_k,
            "vector_store_query_mode": "hybrid"
        })
        return kwargs


class SemanticRetrievalStrategy(RetrievalStrategy):
    """语义检索策略"""

    def __init__(self, top_k: int = 8, enable_query_expansion: bool = True):
        config = RetrievalConfig(
            similarity_top_k=top_k,
            query_transform_mode="hyde" if enable_query_expansion else "default",
            response_mode=ResponseMode.REFINE,
            enable_node_keywords=True
        )
        super().__init__(config)


class KeywordRetrievalStrategy(RetrievalStrategy):
    """关键词检索策略"""

    def __init__(self, top_k: int = 15):
        config = RetrievalConfig(
            similarity_top_k=top_k,
            enable_sparse_embedding=True,
            vector_store_query_mode="simple",
            node_postprocessors=["KeywordPostprocessor"]
        )
        super().__init__(config)


class RetrievalStrategyFactory:
    """检索策略工厂"""

    _strategies = {
        RetrievalMode.STANDARD: StandardRetrievalStrategy,
        RetrievalMode.SIMILARITY: SimilarityRetrievalStrategy,
        RetrievalMode.HYBRID: HybridRetrievalStrategy,
        RetrievalMode.SEMANTIC: SemanticRetrievalStrategy,
        RetrievalMode.KEYWORD: KeywordRetrievalStrategy
    }

    @classmethod
    def create_strategy(cls, mode: RetrievalMode, **kwargs) -> RetrievalStrategy:
        """
        创建检索策略

        Args:
            mode: 检索模式
            **kwargs: 策略参数

        Returns:
            RetrievalStrategy: 检索策略实例
        """
        if mode not in cls._strategies:
            raise ValueError(f"不支持的检索模式: {mode}")

        strategy_class = cls._strategies[mode]
        return strategy_class(**kwargs)

    @classmethod
    def register_strategy(cls, mode: RetrievalMode, strategy_class: type) -> None:
        """
        注册自定义检索策略

        Args:
            mode: 检索模式
            strategy_class: 策略类
        """
        cls._strategies[mode] = strategy_class

    @classmethod
    def get_available_strategies(cls) -> List[RetrievalMode]:
        """
        获取可用的检索策略

        Returns:
            List[RetrievalMode]: 可用的检索模式列表
        """
        return list(cls._strategies.keys())

    @classmethod
    def get_strategy_info(cls) -> Dict[str, str]:
        """
        获取策略信息

        Returns:
            Dict[str, str]: 策略信息字典
        """
        return {
            "standard": "标准向量检索，适合大多数场景",
            "similarity": "基于相似度阈值的精确检索",
            "hybrid": "向量和关键词混合检索，提高召回率",
            "semantic": "语义增强检索，支持查询扩展",
            "keyword": "关键词检索，适合术语密集查询"
        }


def create_default_config(mode: RetrievalMode = RetrievalMode.STANDARD) -> RetrievalConfig:
    """
    创建默认配置

    Args:
        mode: 检索模式

    Returns:
        RetrievalConfig: 默认配置
    """
    strategy_configs = {
        RetrievalMode.STANDARD: RetrievalConfig(),
        RetrievalMode.SIMILARITY: RetrievalConfig(similarity_cutoff=0.7, similarity_top_k=8),
        RetrievalMode.HYBRID: RetrievalConfig(
            similarity_top_k=6,
            enable_sparse_embedding=True,
            vector_store_query_mode="hybrid"
        ),
        RetrievalMode.SEMANTIC: RetrievalConfig(
            similarity_top_k=8,
            query_transform_mode="hyde",
            response_mode=ResponseMode.REFINE
        ),
        RetrievalMode.KEYWORD: RetrievalConfig(
            similarity_top_k=15,
            enable_sparse_embedding=True,
            vector_store_query_mode="simple"
        )
    }

    return strategy_configs.get(mode, RetrievalConfig())